package cn.index.dev.tech.trigger.http;

import cn.index.dev.tech.api.IAiService;
import cn.index.dev.tech.api.model.GenerateStreamRagRequest;
import jakarta.annotation.Resource;
import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.chat.Generation;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.ollama.OllamaChatClient;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.vectorstore.PgVectorStore;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.web.bind.annotation.*;
import org.springframework.http.MediaType;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;


@RestController
@CrossOrigin("*")
@RequestMapping("api/v1/ollama/")
public class OllamaController implements IAiService {

    @Resource
    private OllamaChatClient chatClient;

    @Resource
    private PgVectorStore store;

    @RequestMapping(value = "generate", method = RequestMethod.GET)
    @Override
    public ChatResponse generate(@RequestParam("model") String model, @RequestParam("message") String message) {
        return chatClient.call(new Prompt(message, OllamaOptions.create().withModel(model)));
    }

    @RequestMapping(value = "generate_stream", method = RequestMethod.GET, produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    @Override
    public Flux<ChatResponse> generateStream(@RequestParam("model") String model, @RequestParam("message") String message) {
        // 服务端统一清洗：去除 <think> 标签与无关提示
        return chatClient.stream(new Prompt(message, OllamaOptions.create().withModel(model)))
                .map(this::sanitizeResponse);
    }

    /**
     * 通过流式方式生成基于检索增强生成（RAG）的回答。
     * 使用 POST 请求以更好地处理复杂和长文本输入。
     */
    @RequestMapping(value = "generate_stream_rag", method = RequestMethod.POST, produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    @Override
    public Flux<ChatResponse> generateStreamRag(@RequestBody GenerateStreamRagRequest request) {
        // 1. 定义系统提示词模板
        final String SYSTEM_PROMPT_TEMPLATE = """
        你是一个专业的助手。请仅根据下面提供的【知识库文档】来回答用户的问题。
        
        重要指令：
        1.  你可以在 <think> </think> 标签中进行思考和推理，但最终给用户的答案必须清晰、简洁。
        2.  **最终答案部分不要包含任何 <think> 标签或思考过程。**
        3.  如果知识库中没有相关信息，请在最终答案中明确回答“抱歉，根据我掌握的信息，无法回答您的问题。”
        4.  你的所有回答都必须使用中文。
        
        【知识库文档】:
        {documents}
        """;

        // 2. 在知识库中查询相关文档（适度放宽检索）
        SearchRequest searchRequest = SearchRequest.query(request.message())
                .withTopK(5)
                .withSimilarityThreshold(0.6)
                .withFilterExpression("knowledge == '" + request.ragTag() + "'");

        List<Document> documents = store.similaritySearch(searchRequest);

        // 3. 回退检索：若为空，去掉阈值并提升 topK 再试
        if (documents == null || documents.isEmpty()) {
            SearchRequest fallbackRequest = SearchRequest.query(request.message())
                    .withTopK(8)
                    .withFilterExpression("knowledge == '" + request.ragTag() + "'");
            documents = store.similaritySearch(fallbackRequest);
        }

        // 4. 仍为空则给出友好提示
        if (documents == null || documents.isEmpty()) {
            return Flux.just(new ChatResponse(List.of(new Generation("抱歉，根据我的知识库，没有找到与您问题相关的信息。"))));
        }

        // 5. 文档内容截断，避免系统提示过长
        String documentContext = documents.stream()
                .map(doc -> "文档片段:\n" + (doc.getContent() == null ? "" : (doc.getContent().length() > 1200 ? doc.getContent().substring(0, 1200) : doc.getContent())))
                .collect(Collectors.joining("\n\n---\n\n"));

        // 6. 构建消息
        Message systemMessage = new SystemPromptTemplate(SYSTEM_PROMPT_TEMPLATE)
                .createMessage(Map.of("documents", documentContext));
        List<Message> messages = List.of(systemMessage, new UserMessage(request.message()));

        // 7. 流式调用（移除首事件，直接输出模型结果）
        Prompt prompt = new Prompt(messages, OllamaOptions.create().withModel(request.model()));
        return chatClient.stream(prompt)
                .map(this::sanitizeResponse);
    }

    /**
     * 通过 GET + SSE 的方式支持 RAG 流式，以便前端使用 EventSource。
     */
    @RequestMapping(value = "generate_stream_rag", method = RequestMethod.GET, produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ChatResponse> generateStreamRagGet(@RequestParam("model") String model,
                                                   @RequestParam("ragTag") String ragTag,
                                                   @RequestParam("message") String message) {
        // 1. 定义系统提示词模板
        final String SYSTEM_PROMPT_TEMPLATE = """
        你是一个专业的助手。请仅根据下面提供的【知识库文档】来回答用户的问题。

        重要指令：
        1.  你可以在 <think> </think> 标签中进行思考和推理，但最终给用户的答案必须清晰、简洁。
        2.  **最终答案部分不要包含任何 <think> 标签或思考过程。**
        3.  如果知识库中没有相关信息，请在最终答案中明确回答“抱歉，根据我掌握的信息，无法回答您的问题。”
        4.  你的所有回答都必须使用中文。

        【知识库文档】:
        {documents}
        """;

        // 2. 在知识库中查询相关文档（适度放宽检索）
        SearchRequest searchRequest = SearchRequest.query(message)
                .withTopK(5)
                .withSimilarityThreshold(0.6)
                .withFilterExpression("knowledge == '" + ragTag + "'");
        List<Document> documents = store.similaritySearch(searchRequest);

        // 3. 回退检索：若为空，去掉阈值并提升 topK 再试
        if (documents == null || documents.isEmpty()) {
            SearchRequest fallbackRequest = SearchRequest.query(message)
                    .withTopK(8)
                    .withFilterExpression("knowledge == '" + ragTag + "'");
            documents = store.similaritySearch(fallbackRequest);
        }

        // 4. 仍为空则给出友好提示
        if (documents == null || documents.isEmpty()) {
            return Flux.just(new ChatResponse(List.of(new Generation("抱歉，根据我的知识库，没有找到与您问题相关的信息。"))));
        }

        // 5. 拼接文档上下文（截断）
        String documentContext = documents.stream()
                .map(doc -> "文档片段:\n" + (doc.getContent() == null ? "" : (doc.getContent().length() > 1200 ? doc.getContent().substring(0, 1200) : doc.getContent())))
                .collect(Collectors.joining("\n\n---\n\n"));

        // 6. 构建消息
        Message systemMessage = new SystemPromptTemplate(SYSTEM_PROMPT_TEMPLATE)
                .createMessage(Map.of("documents", documentContext));
        List<Message> messages = List.of(systemMessage, new UserMessage(message));

        // 7. 流式调用（移除首事件，直接输出模型结果）
        Prompt prompt = new Prompt(messages, OllamaOptions.create().withModel(model));
        return chatClient.stream(prompt)
                .map(this::sanitizeResponse);
    }

    /**
     * 统一清洗输出内容：
     * - 去除 <think> ... </think> 片段
     * - 去除以【RAG】开头的任何提示语
     */
    private ChatResponse sanitizeResponse(ChatResponse resp) {
        try {
            List<Generation> gens = resp.getResults();
            if (gens == null || gens.isEmpty()) return resp;
            List<Generation> sanitized = gens.stream().map(g -> {
                String content = g.getOutput().getContent();
                if (content == null) return g;
                // 去除 <think> 标签块
                String c = content.replaceAll("<think>[\\s\\S]*?</think>", "");
                // 去除以【RAG】开头的提示
                c = c.replaceAll("^【RAG】[\\s\\S]*?(\\n|$)", "");
                return new Generation(c);
            }).toList();
            return new ChatResponse(sanitized);
        } catch (Exception e) {
            return resp;
        }
    }

}
